#include "PatchSolutions.hpp"
#include "Misc.hpp"
#include <memory.h>

namespace nkg {

    const uint8_t PatchSolution1::Keyword[0x188] = {
        0xfe, 0xfd, 0xfc, 0xf4, 0xfe, 0xd2, 0xf8, 0xf4, 0xf1, 0xd3, 0xde, 0xc7, 0xdf, 0xd3, 0xd0, 0xfd,
        0x8a, 0xc3, 0x85, 0xf4, 0xf6, 0xe9, 0xfc, 0xfc, 0xf2, 0xf5, 0xfa, 0xf5, 0xf6, 0xe9, 0x81, 0xfb,
        0xfe, 0xfd, 0xfc, 0xf4, 0xf4, 0xdf, 0xf2, 0xf9, 0xf2, 0xe5, 0xf0, 0xf7, 0xc0, 0x89, 0xdd, 0xcb,
        0xf5, 0x87, 0xe6, 0xdd, 0xf4, 0xd9, 0xf8, 0xfb, 0xde, 0xf9, 0xcf, 0xc5, 0x8f, 0x80, 0x80, 0xf3,
        0xc2, 0xd0, 0xe2, 0x8f, 0xfa, 0x8a, 0xdd, 0xf3, 0xd7, 0xdc, 0x86, 0xdc, 0xf0, 0x81, 0xc0, 0xea,
        0xd0, 0xd9, 0xf9, 0xd8, 0xda, 0xf2, 0xd0, 0xfd, 0xc3, 0xf6, 0xf3, 0x82, 0xf2, 0x81, 0xef, 0xf2,
        0xe0, 0xf9, 0xf2, 0xd3, 0x8f, 0xd7, 0xe9, 0xfb, 0xca, 0x86, 0xde, 0xfc, 0xf3, 0xd5, 0xdd, 0xf4,
        0xc7, 0x80, 0xf7, 0xd5, 0xf2, 0xc1, 0xde, 0xcc, 0xc0, 0xc7, 0xf0, 0xd0, 0xd0, 0xd1, 0xd7, 0xcc,
        0xd2, 0x81, 0xc1, 0x83, 0xdd, 0xd5, 0x8a, 0x8f, 0x81, 0xe1, 0xf4, 0xd9, 0xf3, 0xd7, 0xca, 0xef,
        0xf9, 0xdf, 0xe1, 0xee, 0xf0, 0xe9, 0xd1, 0xca, 0xf2, 0xe3, 0xf8, 0xf0, 0x83, 0xde, 0xfb, 0xd7,
        0xf1, 0xc4, 0xfa, 0x85, 0xf2, 0xdd, 0xdd, 0xfd, 0x85, 0x86, 0xc7, 0xf9, 0xc4, 0xc9, 0xf4, 0xf8,
        0xd4, 0xd9, 0xe6, 0xd2, 0xf6, 0xc1, 0xc1, 0xf9, 0xe0, 0xe4, 0xf7, 0xe4, 0xfd, 0xf1, 0xf6, 0xfc,
        0xe1, 0x84, 0xe4, 0xd1, 0xed, 0xfe, 0xdb, 0xe8, 0xdd, 0xe1, 0x85, 0xd0, 0xc5, 0xd2, 0x8a, 0x8e,
        0xd5, 0xdd, 0xe3, 0xdb, 0xd0, 0xe1, 0xd0, 0xf6, 0xc6, 0xee, 0xe6, 0xf7, 0xda, 0xf1, 0xdb, 0xc9,
        0x8b, 0xee, 0xcd, 0xdf, 0xff, 0xe8, 0xdd, 0xca, 0x82, 0xdb, 0xf1, 0x82, 0xc3, 0xed, 0xc9, 0xcc,
        0xc0, 0xf2, 0xd6, 0xdf, 0x83, 0xe9, 0xf3, 0xce, 0xea, 0xfa, 0xdf, 0xf8, 0xd9, 0xff, 0xec, 0x88,
        0xe4, 0xe4, 0xfd, 0x80, 0xc5, 0xce, 0xfa, 0xd2, 0xf4, 0xd8, 0x84, 0xff, 0xe5, 0xf3, 0xcb, 0xc2,
        0xfe, 0xc0, 0xc4, 0xfa, 0xde, 0xdd, 0xd5, 0xc9, 0xc5, 0xd5, 0xdf, 0xe3, 0xdd, 0xc1, 0xcb, 0xdd,
        0xfc, 0xf7, 0x83, 0xf8, 0xda, 0xc1, 0xd4, 0xe3, 0xfe, 0xc2, 0xef, 0xf8, 0xf2, 0xea, 0x8a, 0xd2,
        0xc7, 0xf2, 0xf0, 0xc2, 0xfb, 0x89, 0xdc, 0xeb, 0xd1, 0xf7, 0xcc, 0xe2, 0xd1, 0xfc, 0xd4, 0xce,
        0xea, 0xcd, 0xe4, 0x87, 0xe0, 0xcc, 0x8d, 0xf5, 0xc7, 0x85, 0x87, 0xda, 0xcf, 0xde, 0x89, 0xcd,
        0xe5, 0xfd, 0xe7, 0x83, 0xda, 0xdb, 0xfe, 0xf4, 0x84, 0xec, 0xf6, 0xee, 0xfd, 0xea, 0xf1, 0xf5,
        0xf5, 0xfc, 0xe6, 0xd0, 0x86, 0xdf, 0xc3, 0xe2, 0xe4, 0xd5, 0xd7, 0xe4, 0xe4, 0xce, 0xd4, 0xce,
        0x82, 0xda, 0xc7, 0xda, 0x80, 0xcb, 0xee, 0x8c, 0xd0, 0xde, 0xcd, 0xda, 0xdd, 0xcd, 0xcc, 0xeb,
        0xd2, 0xc3, 0xfc, 0xf2, 0xf6, 0xe9, 0xf8, 0xf8
    };

    PatchSolution1::PatchSolution1(const X64ImageInterpreter& Image) :
        m_Image(Image) {}

    [[nodiscard]]
    bool PatchSolution1::FindPatchOffset() noexcept {
        try {
            void* lpPatch = nullptr;

            for (size_t i = 0; i < m_Image.NumberOfSections(); ++i) {
                auto section = m_Image.ImageSection(i);
                if (section->size >= sizeof(Keyword)) {
                    lpPatch = m_Image.SearchSection(section, [](const void* base, size_t i, size_t size) {
                        if (i + sizeof(Keyword) <= size) {
                            return memcmp(ARL::AddressOffset(base, i), Keyword, sizeof(Keyword)) == 0;
                        } else {
                            return false;
                        }
                    });

                    if (lpPatch) {
                        m_PatchOffset = m_Image.ConvertPtrToOffset(lpPatch);
                        break;
                    }
                }
            }

            if (lpPatch) {
                printf("[+] PatchSolution1 ...... Ready to apply.\n");
                printf("    Keyword offset = +0x%.8x\n", m_PatchOffset.value());
                return true;
            } else {
                throw ARL::Exception(__FILE__, __LINE__, "not found.");
            }
        } catch (...) {
            printf("[-] PatchSolution1 ...... Omitted.\n");
            return false;
        }
    }

    [[nodiscard]]
    bool PatchSolution1::CheckKey(const RSACipher& Cipher) const noexcept {
        try {
            return Cipher.Bits() == 2048;
        } catch (...) {
            return false;
        }
    }

    void PatchSolution1::MakePatch(const RSACipher& Cipher) const {
        if (m_PatchOffset.has_value()) {
            auto szPublicKeyPEM = Cipher.ExportKeyString<RSAKeyType::PublicKey, RSAKeyFormat::PEM>();

            for (auto pos = szPublicKeyPEM.find("-----BEGIN PUBLIC KEY-----"); pos != std::string::npos; pos = szPublicKeyPEM.find("-----BEGIN PUBLIC KEY-----", pos)) {
                szPublicKeyPEM.erase(pos, strlen("-----BEGIN PUBLIC KEY-----"));
            }

            for (auto pos = szPublicKeyPEM.find("-----END PUBLIC KEY-----"); pos != std::string::npos; pos = szPublicKeyPEM.find("-----END PUBLIC KEY-----", pos)) {
                szPublicKeyPEM.erase(pos, strlen("-----END PUBLIC KEY-----"));
            }

            for (auto pos = szPublicKeyPEM.find('\n'); pos != std::string::npos; pos = szPublicKeyPEM.find('\n', pos)) {
                szPublicKeyPEM.erase(pos, 1);
            }
            
            for (size_t i = 0, key = 8; i < szPublicKeyPEM.length(); ++i, --key) {
                if (key == 0) {
                    key = 8;
                }
                szPublicKeyPEM[i] ^= static_cast<char>(0xBB - key);
            }

            while (szPublicKeyPEM.length() < sizeof(Keyword)) {
                szPublicKeyPEM.push_back('\x00');
            }

            auto lpPatch = m_Image.ImageOffset(m_PatchOffset.value());

            puts("**************************************************************");
            puts("*                      PatchSolution1                        *");
            puts("**************************************************************");

            printf("[*] Previous:\n");
            Misc::PrintMemory(lpPatch, sizeof(Keyword), m_Image.ImageBase());
                memcpy(lpPatch, szPublicKeyPEM.data(), sizeof(Keyword));
            printf("[*] After:\n");
            Misc::PrintMemory(lpPatch, sizeof(Keyword), m_Image.ImageBase());
            printf("\n");

        } else {
            throw ARL::Exception(__FILE__, __LINE__, "PatchSolution1: not ready yet.");
        }
    }

}

